#include "StdAfx.h"

int32_t wmain(int32_t nArgc, const wchar_t* pArgv[]) {
	vector<wstring> Args(&pArgv[0], &pArgv[0 + nArgc]);
	shared_ptr<Interface> Intf = make_shared<Interface>(Args);

	if (nArgc < 5) {
		Intf->Log("* Usage: %ws -s <allocation size in hex> -c <allocation count> -a <desired address in hex>\r\n", pArgv[0]);
	}
	else {
		uint32_t dwAllocSize = 0, dwCount = 0;
		uint8_t* pDesiredAddress = nullptr;

		for (vector<wstring>::const_iterator i = Args.begin(); i != Args.end(); ++i) {
			wstring Arg = *i;
			transform(Arg.begin(), Arg.end(), Arg.begin(), ::tolower);

			if (Arg == L"-s") {
				dwAllocSize = wcstoul((*(i + 1)).c_str(), NULL, 0);
			}
			else if (Arg == L"-a") {
#ifdef _WIN64
				pDesiredAddress = reinterpret_cast<uint8_t*>(wcstoull((*(i + 1)).c_str(), NULL, 0));
#else
				pDesiredAddress = reinterpret_cast<uint8_t *>(wcstoul((*(i + 1)).c_str(), NULL, 0));
#endif
			}
			else if (Arg == L"-c") {
				dwCount = _wtoi((*(i + 1)).c_str());
			}
		}

		Intf->Log("... making %d allocations of size 0x%08x\r\n", dwCount, dwAllocSize);

		for (uint32_t dwX = 0; dwX < dwCount; dwX++) {
			uint8_t* pAddress = (uint8_t*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, dwAllocSize);

			if (pAddress == nullptr) {
				Intf->Log("... HeapAlloc failed with error %d (out of memory?)\r\n", GetLastError());
				break;
			}
			else {
				Intf->Log("%d - 0x%p\r\n", dwX, pAddress);

				if (pDesiredAddress != nullptr) {
					if (pDesiredAddress >= pAddress && pDesiredAddress < (pAddress + dwAllocSize)) {
						Intf->Log("! desired address of 0x%p falls within block at 0x%p\r\n", pDesiredAddress, pAddress);
						system("pause");
					}
				}
			}
		}

		system("pause");
	}

	return 0;
}